data {
  int<lower=0> N;         // number of data points
  int<lower=0> Nmunicipalities;        
  int<lower=0> municipality[N];
  int<lower=0> Ymtas[N];
  int<lower=0> Nmtas[N];
  int<lower=0> covid[N]; //
  int<lower=0> Nage_class;         // 
  real age_class_scaled[N];
  int<lower=0> sex[N]; //
  real<lower=0, upper=1> rho_as[N]; // 
}
parameters {
  real log_h0m_raw[Nmunicipalities];
  real<lower=0, upper=1> nu_m[Nmunicipalities];
  real<lower=0> lambda;
  real mu_h;
  real beta_age;
  real beta_sex;
  real<lower=0> sigma_h;
}
  transformed parameters {
    real log_h0m[Nmunicipalities];

    for (m in 1:Nmunicipalities){
      log_h0m[m] = log_h0m_raw[m]*sigma_h + mu_h;
    }
  }
model {
  real lambda_eff[N];
  mu_h ~ normal(-5, 5);
  sigma_h ~ normal(1, 10);
  beta_age ~ normal(0, 2);
  beta_sex ~ normal(0, 2);

  nu_m ~ exponential(lambda);
  //lambda ~ student_t(2, 0, 0.5);
  lambda ~ gamma(0.001, 0.001);

  log_h0m_raw ~ std_normal();

  for (i in 1:N) {
      lambda_eff[i] = Nmtas[i]*(exp(log_h0m[municipality[i]] + beta_age*age_class_scaled[i] + beta_sex*sex[i]) + covid[i]*nu_m[municipality[i]]*rho_as[i]);
    }
  // print(max(lambda_eff));
//  print(Ymtas);

  Ymtas ~ poisson(lambda_eff);
}
// generated quantities {
//  int<lower=0> NbDeath_pred[N];
//  real lambda_eff_pred[N];
//  for (i in 1:N) {
//    lambda_eff_pred[i] = Nmtas[i]*(exp(log_h0m[municipality[i]] + beta_age*age_class_scaled[i] + beta_sex*sex[i]) + covid[i]*nu_m[municipality[i]]*rho_as[i]);
//  }
//   print(max(lambda_eff_pred));
//   NbDeath_pred = poisson_rng(lambda_eff_pred);
// }

